from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
from qwen_vl_utils import process_vision_info
import torch
import util
from PIL import Image

# default: Load the model on the available device(s)
model = Qwen2VLForConditionalGeneration.from_pretrained(
    "hf-models/Qwen2-VL-7B-Instruct-AWQ", torch_dtype="auto", device_map="auto"
)

# We recommend enabling flash_attention_2 for better acceleration and memory saving, especially in multi-image and video scenarios.
# model = Qwen2VLForConditionalGeneration.from_pretrained(
#     "Qwen/Qwen2-VL-2B-Instruct",
#     torch_dtype=torch.bfloat16,
#     attn_implementation="flash_attention_2",
#     device_map="auto",
# )

# default processer
processor = AutoProcessor.from_pretrained("hf-models/hf-models/Qwen2-VL-7B-Instruct-AWQ")


def recognize(image_url, target_text):
    # torch.cuda.empty_cache()
    # The default range for the number of visual tokens per image in the model is 4-16384. You can set min_pixels and max_pixels according to your needs, such as a token count range of 256-1280, to balance speed and memory usage.
    # min_pixels = 256*28*28
    # max_pixels = 1280*28*28
    # processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
    example_json = {"产品名称": "新型冠状病毒(2019-nCoV)抗原检测试剂盒（胶体金法）",  "生产商": "北京乐普诊断科技股份有限公司" }
    prompt = f'''
        ## 角色：
        你是一个高效的智能助手，可以精准、广泛地根据提取目标从图片中识别相应信息。
        ## 任务：
        你现在的任务是从图片中提取相关信息，提取目标包括：{target_text}，各提取目标使用|分隔，请自行拆分，可寻找相似目标，但禁止随意识别不相关信息。
        ## 步骤：
        1. 将提取目标 {target_text} 根据“|”符号进行分隔，并分组管理；
        2. 针对每组的目标，分析相近短语，方便后续识别时更加泛化；
        3. 针对每组的目标，按照后续“要求”进行识别，要考虑“步骤2”中的相近短语；
        4. 按照要求进行结果返回。
        参考示例：
        当提取目标为“产品名称|生产商”时，识别结果为：
        ```json{example_json}```
        ## 要求：
        1. 返回结果使用json格式，每条数据包含一个key-value对，key值为我指定的关键信息，value值为所抽取的结果。
        2. 如果认为OCR识别结果中没有关键信息key，则将value赋值为“未找到相关信息”即可。如果图片模糊，请不要随意猜测内容，将value赋值为“无法识别”即可。 
        3. 直接输出json格式的结果，严禁输出其他内容，包括但不限于：提示语、注释、符号等。
        '''
    util.download_and_show_image(image_url)
    print("image_url", image_url)
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": 'file://'+util.image_path,
                },
                {"type": "text", "text": prompt},
            ],
        }
    ]

    # Preparation for inference
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    )
    inputs = inputs.to("cuda")

    # Inference: Generation of the output
    generated_ids = model.generate(**inputs, max_new_tokens=128)
    generated_ids_trimmed = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    print(output_text[0])
    return output_text[0]

